batch_size = 1 # 批量大小
seq_len = 3 # 样本数量
input_size = 4 # 维度
hidden_size = 2 # 隐藏层(输出)维度
num_layers = 2 # RNN层数

# input.shape = (seq_len, batch_size, input_size)
# h_0.shape = (num_layers, batch_size, hidden_size)
# output.shape = (seq_len, batch_size, hidden_size)
# h_n.shape = (num_layers, batch_size, hidden_size)

# bath_first = True
# input.shape = (batch_size, seq_len, input_size)
# output.shape = (batch_size, seq_len, hidden_size)

import torch

cell = torch.nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)

inputs = torch.randn(seq_len, batch_size, input_size)
hidden = torch.zeros(num_layers, batch_size, hidden_size)

out, hidden = cell(inputs, hidden)

print("Output size: ", out.shape)
print("Output: ", out)
print("Hidden size: ", hidden.shape)
print("Hidden: ", hidden)
